{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "[https://arxiv.org/pdf/2405.04434] Multi-Latent Attention (introduced in DeepSeek-V2)\n",
    "\n",
    "While the attention computation involves less compute than the MLPs in terms of FLOPs for common sequence lengths, \n",
    "many important prompts that users care about are long-sequence, ie. attention handles tens or hundreds of thousands of tokens. \n",
    "These might include putting a novel or paper on in context to get information about it, \n",
    "a Deep Research agent compiling your report with lots of search results in context, and more. \n",
    "\n",
    "The primary bottleneck in these long-sequence workloads is the memory required to store the KV cache, \n",
    "which is linear in S, and is roughly O(bsd) where b is batch size, s is seqlen, d is mode's hidden dim. \n",
    "MLA is a way to compute attention *in a smaller latent space* which means *we never need to materialize keys and values \n",
    "in the computation at all*. The core idea is to just store latent vectors c_kv in the cache, so the cache size now \n",
    "goes from [b, nh, s, hd] to now [b, nh, s, d_c_h] where (d_c = d_c_h * nh) << (d_model = d_h * nh). \n",
    "\n",
    "Queries are also projected to a small latent space to help save memory during training. Some important\n",
    "conceptual notes \n",
    "    - The paper introduces a lot of matrices re up/down projections. self.to_c_kv \"absorbs\" most of them in that \n",
    "    we don't explicitly need all of those, they are learned when we learn to_c_kv. \n",
    "    - Attn itself is computed in latent space, and this is why the full K matrix never appears. \n",
    "\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✓ MLA causal: full‑sequence == incremental decode\n"
     ]
    }
   ],
   "source": [
    "### minimal implementation, ignoring decoupled RoPE for simplicity \n",
    "import torch, torch.nn as nn, torch.nn.functional as F\n",
    "import math\n",
    "from tqdm import tqdm \n",
    "import time\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "class MLAProjections(nn.Module): \n",
    "    def __init__(self, d_model: int = 512, d_c_total: int = 64): \n",
    "        super().__init__()\n",
    "\n",
    "        self.to_q_lat = nn.Linear(d_model, d_c_total, bias=False)\n",
    "        self.to_c_kv = nn.Linear(d_model, d_c_total, bias=False)\n",
    "        self.to_v = nn.Linear(d_c_total, d_model, bias=False)\n",
    "        self.wo = nn.Linear(d_model, d_model, bias=False)\n",
    "\n",
    "\n",
    "class MLA(nn.Module): \n",
    "    def __init__(self, d_model: int = 512, d_c_total: int = 64, n_heads: int = 16, causal: bool = True):\n",
    "        super().__init__()\n",
    "\n",
    "        assert d_model % n_heads == 0, \"Error: n_heads must divide d_model for \\\n",
    "            each head to be of the same size!\"\n",
    "        assert d_c_total % n_heads == 0, \"Error: n_heads must divide d_c_total!\"\n",
    "    \n",
    "        self.d_model = d_model \n",
    "        self.d_c_total = d_c_total \n",
    "        self.n_heads = n_heads \n",
    "        self.d_h = int(self.d_model/n_heads)\n",
    "        self.d_c_h = int(d_c_total/n_heads)\n",
    "        self.causal = causal \n",
    "        self.scale = 1.0 / math.sqrt(self.d_c_h)\n",
    "\n",
    "        self.projections = MLAProjections(d_model, d_c_total)\n",
    "\n",
    "        # need to have a KV cache since that's where the gains are! \n",
    "        # see language-models/kv_cache.ipynb if you haven't already \n",
    "        self.clear_cache()\n",
    "\n",
    "    def clear_cache(self):\n",
    "        self.register_buffer(\"c_cache\", None, persistent=False)\n",
    "        self.cache_pos = 0\n",
    "\n",
    "    def _apply_causal_mask(self, logits: torch.Tensor, s: int, L: int, device: torch.device):\n",
    "        \"\"\"Apply causal mask to attention logits\"\"\"\n",
    "        q_idx = torch.arange(s, device=device)\n",
    "        k_idx = torch.arange(L, device=device)\n",
    "        mask = (k_idx[None, :] > q_idx[:, None])                      # [s,k]\n",
    "        logits.masked_fill_(mask[None, None, ...], float('-inf'))\n",
    "\n",
    "    def forward(self, x: torch.Tensor, decode: bool = False): # [b, s, d_model] -> [b, s, d_model]\n",
    "        b, s, d = x.shape \n",
    "        device = x.device \n",
    "\n",
    "        # get q, kv (via c) in latent space \n",
    "        q_lat = self.projections.to_q_lat(x) # [b, s, d_c]\n",
    "        c_kv = self.projections.to_c_kv(x)   # [b, s, d_c]\n",
    "\n",
    "        # reshape both [b, s, d_c] ->  [b, nh, s, d_c_h]\n",
    "        q_lat = q_lat.view(b, s, self.n_heads, self.d_c_h).transpose(1, 2) \n",
    "        c_kv = c_kv.view(b, s, self.n_heads, self.d_c_h).transpose(1, 2) ## <-- the key new object in MLA compared to vanilla mhsa\n",
    "\n",
    "        # hit/update cache if decode \n",
    "        if decode: \n",
    "            if self.cache_pos == 0: # this is the only place we use .cache_pos\n",
    "                self.c_cache = c_kv\n",
    "                self.cache_pos += s                   \n",
    "                k_len = s\n",
    "            else: \n",
    "                c_kv = torch.cat([self.c_cache, c_kv], dim=2) \n",
    "                self.c_cache = c_kv \n",
    "                self.cache_pos += s\n",
    "                k_len = self.c_cache.size(2)\n",
    "        else:\n",
    "            self.clear_cache()\n",
    "            k_len = s\n",
    "\n",
    "        # do attn in latent space \n",
    "        attn_logits = torch.einsum('bhqd,bhkd->bhqk', q_lat, c_kv) * self.scale\n",
    "        \n",
    "        if self.causal and not decode: \n",
    "            self._apply_causal_mask(attn_logits, s, s, device)\n",
    "            \n",
    "        attn = F.softmax(attn_logits, dim=-1)\n",
    "\n",
    "        # project c_kv to values and do A@v \n",
    "        v = self.projections.to_v(c_kv.transpose(1, 2).reshape(b, k_len, self.d_c_total))\n",
    "        v = v.view(b, k_len, self.n_heads, self.d_h).transpose(1, 2)\n",
    "\n",
    "        # return wo(out)\n",
    "        out = torch.einsum('bhqk,bhkd->bhqd', attn, v)\n",
    "        out = out.transpose(1, 2).reshape(b, s, -1)\n",
    "        \n",
    "        return self.projections.wo(out)\n",
    "\n",
    "\n",
    "# sanity check cache and decoding is working correctly\n",
    "if __name__ == \"__main__\":\n",
    "    torch.manual_seed(0)\n",
    "    mla = MLA(d_model=256, n_heads=4, d_c_total=64)\n",
    "    x   = torch.randn(2, 10, 256) # b, s, d\n",
    "\n",
    "    # full‑sequence (masked)\n",
    "    y_full = mla(x)                              \n",
    "\n",
    "    # incremental decode\n",
    "    mla.clear_cache()\n",
    "    y_inc = torch.cat([mla(x[:, t:t+1, :], decode=True) for t in range(10)], dim=1)\n",
    "\n",
    "    # they should now be identical\n",
    "    torch.testing.assert_close(y_full, y_inc, atol=1e-5, rtol=0)\n",
    "    print(\"✓ MLA causal: full‑sequence == incremental decode\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vanilla MHSA max sequence length: 11920\n",
      "MLA max sequence length: 14900\n"
     ]
    }
   ],
   "source": [
    "# baseline we'll compare MLA to, ignore this cell it's just defining a baseline \n",
    "class VanillaMHSA(nn.Module):\n",
    "    def __init__(self, d_model, n_heads):\n",
    "        super().__init__()\n",
    "        self.d_model = d_model\n",
    "        self.n_heads = n_heads\n",
    "        self.d_h = d_model // n_heads\n",
    "        self.scale = self.d_h ** -0.5\n",
    "        self.q_proj = nn.Linear(d_model, d_model)\n",
    "        self.k_proj = nn.Linear(d_model, d_model)\n",
    "        self.v_proj = nn.Linear(d_model, d_model)\n",
    "        self.out_proj = nn.Linear(d_model, d_model)\n",
    "        self.clear_cache()\n",
    "\n",
    "    def clear_cache(self):\n",
    "        self.k_cache = None\n",
    "        self.v_cache = None\n",
    "\n",
    "    def forward(self, x, decode=False):\n",
    "        b, s, _ = x.shape\n",
    "        # project q, k, v\n",
    "        q = self.q_proj(x).view(b, s, self.n_heads, self.d_h).transpose(1, 2)\n",
    "        k = self.k_proj(x).view(b, s, self.n_heads, self.d_h).transpose(1, 2)\n",
    "        v = self.v_proj(x).view(b, s, self.n_heads, self.d_h).transpose(1, 2)\n",
    "\n",
    "        if decode:\n",
    "            # incremental decode: append new k, v to cache\n",
    "            if self.k_cache is None:\n",
    "                self.k_cache = k\n",
    "                self.v_cache = v\n",
    "            else:\n",
    "                self.k_cache = torch.cat([self.k_cache, k], dim=2)\n",
    "                self.v_cache = torch.cat([self.v_cache, v], dim=2)\n",
    "            k_all, v_all = self.k_cache, self.v_cache\n",
    "        else:\n",
    "            k_all, v_all = k, v\n",
    "\n",
    "        # scaled dot-product\n",
    "        attn_scores = torch.einsum('bhqd,bhkd->bhqk', q, k_all) * self.scale\n",
    "        if decode:\n",
    "            # at decode time we only need last query row\n",
    "            attn_scores = attn_scores[:, :, -1:, :]\n",
    "            attn = F.softmax(attn_scores, dim=-1)\n",
    "        else:\n",
    "            # full-sequence causal mask\n",
    "            mask = torch.tril(torch.ones(s, s, device=x.device)).view(1, 1, s, s)\n",
    "            attn = F.softmax(attn_scores.masked_fill(mask == 0, float('-inf')), dim=-1)\n",
    "\n",
    "        out = torch.einsum('bhqk,bhkd->bhqd', attn, v_all)\n",
    "        out = out.transpose(1, 2).reshape(b, s, self.d_model)\n",
    "        return self.out_proj(out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vanilla MHSA max sequence length: 11920\n",
      "MLA max sequence length: 14900\n"
     ]
    }
   ],
   "source": [
    "# here we see the max seqlen both mhsa and mla can decode to \n",
    "# the upshot of MLA is nominally a smaller KV cache, by roughly a factor of d_c/d_model\n",
    "# so MHSA should OOM earlier than MLA, and we see that indeed it does \n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "d_model, n_heads, d_c_total = 2048, 2048//64, 256\n",
    "\n",
    "mla = MLA(d_model=d_model, n_heads=n_heads, d_c_total=d_c_total).to(device).eval()\n",
    "vanilla = VanillaMHSA(d_model=d_model, n_heads=n_heads).to(device).eval()\n",
    "\n",
    "def try_decode(model, seqlen):\n",
    "    seqlen = int(seqlen)\n",
    "    try:\n",
    "        model.clear_cache()\n",
    "        # create input of the specified sequence length\n",
    "        x = torch.randn(1, seqlen, d_model, device=device)\n",
    "        \n",
    "        # process the entire input at once to build the kv cache\n",
    "        model(x, decode=False)\n",
    "        \n",
    "        # then decode 20 additional tokens\n",
    "        gen_x = torch.randn(1, 1, d_model, device=device)\n",
    "        for _ in range(20):\n",
    "            model(gen_x, decode=True)\n",
    "            \n",
    "        return True\n",
    "    except RuntimeError as e:\n",
    "        if \"out of memory\" in str(e):\n",
    "            return False\n",
    "        raise e\n",
    "    except KeyboardInterrupt:\n",
    "        print(\"Operation interrupted by user\")\n",
    "        return False\n",
    "\n",
    "# try increasingly large sequence lengths until oom\n",
    "seqlen = 1024\n",
    "vanilla_max_seqlen = 0\n",
    "mla_max_seqlen = 0\n",
    "try:\n",
    "    # test vanilla mhsa\n",
    "    while True:\n",
    "        if try_decode(vanilla, int(seqlen)):\n",
    "            vanilla_max_seqlen = int(seqlen)\n",
    "            seqlen *= 1.25\n",
    "        else:\n",
    "            break\n",
    "    \n",
    "    print(f\"Vanilla MHSA max sequence length: {vanilla_max_seqlen}\")\n",
    "    \n",
    "    # test mla\n",
    "    seqlen = vanilla_max_seqlen\n",
    "    while True:\n",
    "        if try_decode(mla, int(seqlen)):\n",
    "            mla_max_seqlen = int(seqlen)\n",
    "            seqlen *= 1.25\n",
    "        else:\n",
    "            break\n",
    "    \n",
    "    print(f\"MLA max sequence length: {mla_max_seqlen}\")\n",
    "    \n",
    "except KeyboardInterrupt:\n",
    "    print(\"\\nBenchmark interrupted by user\")\n",
    "    if vanilla_max_seqlen > 0:\n",
    "        print(f\"Vanilla MHSA max sequence length: {vanilla_max_seqlen}\")\n",
    "    if mla_max_seqlen > 0:\n",
    "        print(f\"MLA max sequence length: {mla_max_seqlen}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "envi",
   "language": "python",
   "name": "lingua_env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
